package com.atguigu.distributed.lock.leecode;

import java.util.Comparator;
import java.util.PriorityQueue;

/**
 * 输入 n 个整数，找出其中最小的 k 个数。例如输入4、5、1、6、2、7、3、8 这8个数字，则最小的4个数字是1、2、3、4
 */
public class GetKLittleNum {
    public static void main(String[] args) {
        int[] input = {1,2,3,4,5,6,7,8,9};
        int[] ouput = new int[5];
        //getLeastNumbers(input,ouput);
        Integer[] leastNumbers = getLeastNumbers(input, 5);
        for (int i = 0; i < leastNumbers.length; i++) {
            System.out.println(leastNumbers[i]);
        }

    }

    /**
     * 此方法适合数据量一般的，不适用于海量数据的最小的k个数
     * @param input
     * @param output
     */
    public static void getLeastNumbers(int[] input, int[] output) {
        if (input == null || output == null || output.length <= 0 || input.length < output.length) {
            throw new IllegalArgumentException("Invalid args");
        }

        int start = 0;
        int end = input.length - 1;
        int index = partition(input, start, end); //切分后左子数组的长度
        int target = output.length - 1; //K-1

        //若切分后左子数组长度不等于K
        while (index != target) {
            //若切分后左子数组长度小于K，那么继续切分右子数组，否则继续切分左子数组
            if (index < target) {
                start = index + 1;
            } else {
                end = index - 1;
            }
            index = partition(input, start, end);
        }

        System.arraycopy(input, 0, output, 0, output.length);
    }

    private static int partition(int[] arr, int left, int right) {
        int le = left+1;
        int ge = right;
        while (true){
            while (le<=ge&&arr[le]<arr[left]){
                le++;
            }
            while (le<=ge&&arr[ge]>arr[left]){
                ge--;
            }
            if(le>=ge){
                break;
            }
            swap(arr,le,ge);
            le++;
            ge--;

        }
        swap(arr,left,ge);
        return ge;

    }

    private static void swap(int[] arr, int i, int j) {
        int tmp = arr[i];
        arr[i] = arr[j];
        arr[j] = tmp;
    }

    /**
     * 此方法适用于海量数据中最小的k个数
     * @param nums
     * @param k
     * @return
     */
    public static Integer[] getLeastNumbers(int[] nums,int k){
        //默认自然排序，需手动转为降序
        PriorityQueue<Integer> maxQueue = new PriorityQueue<>(k, new Comparator<Integer>() {
            @Override
            public int compare(Integer o1, Integer o2) {
                if(o1>o2){
                    return -1;
                }else if(o1<o2){
                    return 1;
                }
                return 0;
            }
        });
        for(int num:nums){
            if(maxQueue.size()<k||num<maxQueue.peek()){
                //插入元素
                maxQueue.offer(num);
            }
            if(maxQueue.size()>k){
                maxQueue.poll();
            }
        }
        return maxQueue.toArray(new Integer[0]);

    }

    /**
     * 此方法适用于海量数据中最大的k个数
     * @param nums
     * @param k
     * @return
     */
    public Integer[] getLargestNumber(int[] nums,int k){
        PriorityQueue<Integer> minQueue = new PriorityQueue<>(k);//默认自然排序
        for(int num:nums){
            if(minQueue.size()<k||num>minQueue.peek()){
                minQueue.offer(num);
            }
            if(minQueue.size()>k){
                //删除队列头部
                minQueue.poll();
            }
        }
        return minQueue.toArray(new Integer[0]);
    }

}
